
from struct import pack, unpack, calcsize
from mathutils import Matrix, Vector, Euler, Quaternion

from math import *
from collections import OrderedDict

global file
file = None

class Bone(object):
    __slots__ = (
        "name",  # bone name, used as key in many places
        "parent",  # decorated bone parent, set in a later loop
        "rest_bone",  # blender armature bone
        "pose_bone",  # blender pose bone
        "pose_mat",  # blender pose matrix
        "rest_arm_mat",  # blender rest matrix (armature space)
        "rest_local_mat",  # blender rest batrix (local space)
        "pose_imat",  # pose_mat inverted
        "rest_arm_imat",  # rest_arm_mat inverted
        "rest_local_imat",  # rest_local_mat inverted
        "is_tail_head",  # is the bone disconnected to the parent bone?
        "head_location",  # relative bone head location
        "tail_location",  # relative bone tail location
        "children",
        "mandatory",
        "frame_values",
        "use_loc",
        "use_rot",
        "adj_mat"
    )

    def __repr__(self):
        if self.parent:
            return "[\"%s\" child on \"%s\"]\n" % (self.name, self.parent.name)
        else:
            return "[\"%s\" root bone]\n" % (self.name)

    def __init__(self, armature, bone_name, adj_mat):
        self.name = bone_name
        self.adj_mat = adj_mat
        self.rest_bone = armature.armature.data.bones[bone_name]
        self.pose_bone = armature.armature.pose.bones[bone_name]
        
        self.rest_local_mat = self.rest_bone.matrix
        if (bone_name == "mPelvis"):
            self.rest_local_mat = self.adj_mat.to_3x3() @ self.rest_local_mat
        self.rest_local_imat = self.rest_local_mat.inverted()

        self.pose_mat = self.adj_mat @ self.pose_bone.matrix
        self.pose_imat = self.pose_mat.inverted()

        self.rest_arm_mat = self.adj_mat @ self.rest_bone.matrix_local
        self.rest_arm_imat = self.rest_arm_mat.inverted()
        
        self.parent = None
        self.is_tail_head = (self.rest_bone.use_connect and self.rest_bone.parent)
        
        self.children = []
        self.frame_values = []

        self.mandatory = False
        self.use_rot = False
        self.use_loc = False
        
    def adjust_hierarchy_up(self, armature):
        if self.rest_bone.parent == None: return
        self.parent = armature.bones[self.rest_bone.parent.name]
        self.parent.children.append(self)

    def adjust_hierarchy_down(self, armature):
        self.head_location = self.adj_mat @ self.rest_bone.head_local
        self.tail_location = self.adj_mat @ self.rest_bone.tail_local

        self.children.sort(key=lambda x: x.name)

        for bone in self.children:
            bone.adjust_hierarchy_down(armature)

    def set_mandatory(self):
        if (not self.mandatory) and (self.parent != None):
            self.parent.set_mandatory()
        self.mandatory = True

    def recursive_update_posedata(self):
        self.pose_mat = self.adj_mat @ self.pose_bone.matrix
        self.pose_imat = self.pose_mat.inverted()
        for bone in self.children: bone.recursive_update_posedata()
        
    def recursive_compute_frame(self):
        trans = Matrix.Translation(self.head_location)
        itrans = Matrix.Translation(-self.head_location)

        if self.parent:
            parent_pose_mat = self.parent.pose_mat
            parent_pose_imat = self.parent.pose_imat
            parent_rest_arm_mat = self.parent.rest_arm_mat
            parent_head_location = self.parent.head_location
            head_location = self.head_location
        else:
            parent_pose_mat = Matrix()
            parent_pose_imat = Matrix()
            parent_rest_arm_mat = Matrix()
            parent_head_location = Vector((0,0,0))
            head_location = self.adj_mat @ self.rest_bone.head
            parent_head_location = head_location
        
        from ossl.util import fmt

        pose_mat = self.pose_mat
        rest_arm_imat = self.rest_arm_imat
        mat_final = parent_rest_arm_mat @ parent_pose_imat @ pose_mat @ rest_arm_imat
        mat_final = itrans @ mat_final @ trans
        loc = mat_final.to_translation() + (head_location - parent_head_location)
        rot = mat_final.to_3x3().inverted()

#        import bpy
#        def pp(d):
#            return ### because of debug output only
#            f = self.name == "mPelvis" and bpy.context.scene.frame_current == 0
#            file = open("C:\\Blender\\_addons\\dan.txt", "w" if f else "a", encoding="utf-8")
#            for n in d: file.write(n + ": " + fmt(d[n]) + "\n")
#            file.write("\n")
#            file.close()

#        if bpy.context.scene.frame_current in [0,36] and \
#            self.name in ["mPelvis","PELVIS","mCollarLeft","mShoulderLeft"]:
#            pp({"bone": self.name,
#                "frame": bpy.context.scene.frame_current,
#                "parent.pose_mat": parent_pose_mat,
#                "parent.head_location": parent_head_location,
#                "parent.rest_arm_mat": parent_rest_arm_mat,
#                "self.pose_mat": pose_mat,
#                "self.head_location": head_location,
#                "self.rest_arm_mat": self.rest_arm_mat,
#                "mat_final": mat_final,
#                "loc": loc,
#                "rot": rot,
#                }) 
        self.frame_values.append((loc, rot))

        for bone in self.children: bone.recursive_compute_frame()


import ossl.util as util

def register():
    print("registering: OSSL Type Bone")
    print("registered: OSSL Type Bone")

def unregister():
    print("unregistering: OSSL Type Bone")
    print("unregistered: OSSL Type Bone")

if __name__ == "__main__":
    util.reset_addon()
